from data_preprocessor import *


def test_cleaning_xml():
    xml_path = '/home/tang/workspace/quechao/test_tmp/00004_.xml'
    label_filtering_rule = ['gerber', 'nric', 'Nestle']
    old_label_list = ['gerber12', 'nric3', 'Nestle15']
    valid, label_list = cleaning_xml(xml_path, label_filtering_rule, old_label_list)
    assert valid
    assert label_list != old_label_list


def test_cleaning_dataset():
    dataset_dir = '/home/tang/workspace/quechao/test_tmp'  # 该目录中有三张图片，两张有对应的xml
    label_filtering_rule = ['gerber', 'nric']
    ret = cleaning_dataset(dataset_dir, label_filtering_rule)
    print(ret)
    # print(ret['label_list'])
    assert ret is not None
    assert len(ret) != 0
    assert len(ret['label_list']) != 0
    assert ret['annotation_nums'] == 2
    assert ret['image_nums'] == 3


def test_create_label_map():
    label_list = ['nric9', 'gerber22', 'nric3', 'gerber28', 'nric7', 'gerber21', 'nric16', 'nric14', 'gerber32',
                  'nric12', 'gerber25', 'gerber23']
    output_file_path = 'test_label_map.pbtxt'
    cnt = create_label_map(label_list, output_file_path)
    assert cnt == len(label_list)


def test_create_tfrecordes():
    label_map_path = 'test_label_map.pbtxt'
    dataset_dir = '/home/tang/workspace/quechao/test_tmp'
    output_file_path = 'test.tfrecords'
    create_tfrecords(dataset_dir, label_map_path, output_file_path)
